import os
from pprint import pprint

from dotenv import load_dotenv
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, ConfigurableFieldSpec
from langchain_openai import AzureChatOpenAI
from langchain_community.chat_message_histories import ChatMessageHistory

load_dotenv()

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You're an assistant who's good at {ability}. Respond in 20 words or fewer"),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{input}")
    ]
)

# model = ChatOpenAI(model="gpt-4")
model = AzureChatOpenAI(
    # openai_api_key=
    # openai_api_base=os.getenv("AZURE_OPENAI_ENDPOINT"),
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    azure_deployment=os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME"),
    api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
    temperature=0.7
)

# 字典作为输入，并返回消息作为输出
runnable = prompt | model
pprint(runnable)

"""
用户ID，会话ID
"""
## 字典--- 内存;;也可以存到数据库中。
store = {}


def get_session_history(user_id: str, conversation_id: str) -> BaseChatMessageHistory:
    if (user_id, conversation_id) not in store:
        store[user_id, conversation_id] = ChatMessageHistory()
    return store[user_id, conversation_id]


## RunnableWithMessageHistory 带上历史消息
with_message_history = RunnableWithMessageHistory(
    runnable=runnable,
    get_session_history=get_session_history,
    input_messages_key="input",
    history_messages_key="history",
    history_factory_config=[
        ConfigurableFieldSpec(
            id="user_id",
            annotation=str,
            name="User ID",
            description="用户的唯一标识符",
            default="",
            is_shared=True
        ),
        ConfigurableFieldSpec(
            id="conversation_id",
            annotation=str,
            name="Conversation ID",
            description="对话的唯一标识符",
            default="",
            is_shared=True
        ),
    ],
)

response = with_message_history.invoke(
    input={"ability": "math", "input": "余弦是什么意思"},
    config={"configurable": {"user_id": "123", "conversation_id": "1"}},
)
print(response)
## 输出：content='余弦是三角函数之一，表示一个角的邻边长度与斜边长度之比，常记为cos...


##记住，
response = with_message_history.invoke(
    input={"ability": "math", "input": "什么"},
    config={"configurable": {"user_id": "123", "conversation_id": "1"}},
)
print(response)
## content='余弦是三角函数，定义为直角三角形中邻边长度除以斜边长度，符号为cos。

##新的：conversation_id--> 不记得了
response = with_message_history.invoke(
    input={"ability": "math", "input": "什么"},
    config={"configurable": {"user_id": "123", "conversation_id": "2"}},
)
print(response)
## 输出：content='你好！请问有什么数学问题可以帮您解决的吗？
